from torch import nn
from torch.autograd import Variable
import torch
import math
import unittest

"""
该代码实现了  ​Transformer 模型中的位置编码 （ Positional Encoding） ​， 用于为序列数据（如文本、时间序列）注入位置信息。

​核心作用​：
Transformer 模型依赖自注意力机制，但自注意力本身无法感知序列中元素的位置顺序。位置编码通过为每个位置生成唯一的高维向量，
将其叠加到输入嵌入中，使模型能够区分不同位置的元素。



"""

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_seq_len=80):
        super(PositionalEncoding, self).__init__()
        super().__init__()
        self.d_model = d_model


        """
        代码基于论文《Attention Is All You Need》的正弦余弦位置编码方法：
        [2,100002,2](@ref)其中：
        pos 为位置索引 （ 0 到 max_seq_len-1 ） 。
        i 为编码维度的偶数索引 （ 0, 2, 4...）。
        d_model 为编码维度（如 512 ） 。
        """
        pe = torch.zeros(max_seq_len, d_model)
        for pos in range(max_seq_len):
            for i in range(0, d_model, 2):
                pe[pos, i] = math.sin(pos / (10000 ** ((2 * i)/d_model)))
                pe[pos, i + 1] = math.cos(pos / (10000 ** ((2 * (i + 1))/d_model)))
        
        #生成一个形状为 (max_seq_len, d_model) 的位置编码矩阵 pe，并通过 unsqueeze(0) 扩展为 (1, max_seq_len, d_model)，便于广播到不同长度的输入序列。
        pe=pe.unsqueeze(0)
        #将 pe 注册为模型的 buffer，使其在模型保存和加载时被保留，但不会参与梯度计算。
        self.register_buffer('pe', pe)

    def forward(self, x):
        #对输入 x（形状为 (batch_size, seq_len, d_model)）进行缩放，防止后续加法导致数值不稳定。
        x=x*math.sqrt(self.d_model)
        seq_len=x.size(1)
        #将位置编码 pe 的前 seq_len 列（对应当前输入序列长度）与缩放后的输入 x 相加，形成位置感知的输入。
        #通过 requires_grad=False 禁用位置编码的梯度计算，确保其固定不变。
        #使用 .cuda() 将位置编码移动到 GPU 上，与输入 x 保持设备一致。
        #x=x+Variable(self.pe[:, :seq_len],requires_grad=False)  #.cuda()
        #Variable 在新版 PyTorch 中已弃用，建议直接使用张量的 requires_grad 属性：
        x = x + self.pe[:, :seq_len].to(x.device)
        return x


class TestPositionalEncoding(unittest.TestCase):
    def setUp(self):
        # 初始化测试参数
        self.d_model = 16
        self.max_seq_len = 10
        self.batch_size = 2
        self.seq_len = 5
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def test_initialization(self):
        """测试初始化参数和位置编码矩阵的形状"""
        pe_layer = PositionalEncoding(self.d_model, self.max_seq_len).to(self.device)
        
        # 检查 d_model 属性是否正确
        self.assertEqual(pe_layer.d_model, self.d_model)
        
        # 检查位置编码矩阵形状
        self.assertEqual(pe_layer.pe.shape, (1, self.max_seq_len, self.d_model))
        
        # 检查数值是否有限（无 NaN/Inf）
        self.assertTrue(torch.all(torch.isfinite(pe_layer.pe)))

    def test_positional_encoding_values(self):
        """测试位置编码的数值是否符合公式"""
        pe_layer = PositionalEncoding(self.d_model, self.max_seq_len).to(self.device)
        pe = pe_layer.pe.squeeze(0)  # 去除 batch 维度
        
        # 验证第一个位置 (pos=0) 的编码值
        pos = 0
        for i in range(0, self.d_model, 2):
            expected_sin = math.sin(pos / (10000 ** ((2 * i)/self.d_model)))
            expected_cos = math.cos(pos / (10000 ** ((2 * (i + 1))/self.d_model)))
            self.assertAlmostEqual(pe[pos, i].item(), expected_sin, places=5)
            self.assertAlmostEqual(pe[pos, i + 1].item(), expected_cos, places=5)

    def test_forward(self):
        """测试前向传播的输出形状和数值"""
        pe_layer = PositionalEncoding(self.d_model, self.max_seq_len).to(self.device)
        x = torch.randn(self.batch_size, self.seq_len, self.d_model).to(self.device)
        
        # 前向传播
        output = pe_layer(x)
        
        # 检查输出形状
        self.assertEqual(output.shape, x.shape)
        
        # 检查位置编码是否叠加
        # 手动计算期望的位置编码部分
        expected_pe = pe_layer.pe[:, :self.seq_len]
        # 由于浮点精度问题，使用张量近似比较
        self.assertTrue(torch.allclose(output - x, expected_pe.to(x.device), atol=1e-5))

    def test_device_compatibility(self):
        """测试设备兼容性（CPU/GPU）"""
        pe_layer_cpu = PositionalEncoding(self.d_model, self.max_seq_len).to("cpu")
        x_cpu = torch.randn(self.batch_size, self.seq_len, self.d_model).to("cpu")
        output_cpu = pe_layer_cpu(x_cpu)
        
        pe_layer_gpu = PositionalEncoding(self.d_model, self.max_seq_len).to("cuda")
        x_gpu = x_cpu.to("cuda")
        output_gpu = pe_layer_gpu(x_gpu)
        
        # 检查设备一致性
        self.assertEqual(output_cpu.device, x_cpu.device)
        self.assertEqual(output_gpu.device, x_gpu.device)

    def test_grad_check(self):
        """测试梯度不可训练性"""
        pe_layer = PositionalEncoding(self.d_model, self.max_seq_len).to(self.device)
        x = torch.randn(self.batch_size, self.seq_len, self.d_model, requires_grad=True).to(self.device)
        
        # 前向传播
        output = pe_layer(x)
        
        # 反向传播检查梯度
        output.sum().backward()
        # 检查输入 x 的梯度是否存在
        self.assertIsNotNone(x.grad)
        # 检查位置编码的梯度是否为 None（不可训练）
        self.assertIsNone(pe_layer.pe.grad)

if __name__ == '__main__':
    # 当前安装的 PyTorch 库的版本
    print(torch.__version__)
    # 检查 CUDA 是否可用，即你的系统有 NVIDIA 的 GPU
    print(torch.cuda.is_available())
    unittest.main()